"""
1.	利用深度学习框架，建立由卷积神经网络Resnet34构造的特征提取主干模型，和后端分类网络，进行猫狗图像识别，并进行模型评估，
按下面要求完成相应代码（56分）
"""
import torch as pt
from torch.utils.data import TensorDataset, DataLoader
import torchvision as ptv
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.metrics import precision_score, recall_score, f1_score, confusion_matrix
import seaborn as sns
import os
import cv2 as cv

# 随机种子
np.random.seed(1)
pt.random.manual_seed(1)

# 超参数
VER_INT = 1
VER = f'v{VER_INT}'
SIZE = 224
BATCH_SIZE = 4
N_EPOCHS = 2
ALPHA = 1e-3
BASE_DIR, FILE_NAME = os.path.split(__file__)
dir = 'data/zoo'
IMG_DIR = os.path.join(BASE_DIR, dir)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)

# GPU / CPU
device = 'cuda:0' if pt.cuda.is_available() else 'cpu'
print('device', device)
device = pt.device(device)

# ①	读入猫狗图像数据集，进行适当训练集、验证集、测试集比例划分
x, y = [], []
yi = 0
idx2label, label2idx = {}, {}
for sub_dir_name in os.listdir(IMG_DIR):
    sub_dir_path = os.path.join(IMG_DIR, sub_dir_name)
    idx2label[yi] = sub_dir_name
    label2idx[sub_dir_name] = yi
    for file_name in os.listdir(sub_dir_path):
        file_path = os.path.join(sub_dir_path, file_name)
        img = cv.imread(file_path, cv.IMREAD_COLOR)
        img = cv.resize(img, (SIZE, SIZE), interpolation=cv.INTER_CUBIC)
        x.append(img)
        y.append(yi)
    yi += 1
print(idx2label)
n_cls = len(idx2label)
x = np.float32(x) / 255.
x = np.transpose(x, (0, 3, 1, 2))
y = np.int64(y)
print('x', x.shape)
print('y', y.shape)
x = pt.Tensor(x)
y = pt.Tensor(y)
x_train, x_val_test, y_train, y_val_test = train_test_split(x, y, train_size=0.8, random_state=VER_INT, shuffle=True)
x_val, x_test, y_val, y_test = train_test_split(x_val_test, y_val_test, train_size=0.5, random_state=VER_INT, shuffle=True)
print('x_train', x_train.size())
print('x_val', x_val.size())
print('x_test', x_test.size())
print('y_train', y_train.size())
print('y_val', y_val.size())
print('y_test', y_test.size())
ds_train = TensorDataset(x_train, y_train)
ds_val = TensorDataset(x_val, y_val)
ds_test = TensorDataset(x_test, y_test)
dl_train = DataLoader(ds_train, BATCH_SIZE)
dl_val = DataLoader(ds_val, BATCH_SIZE)
dl_test = DataLoader(ds_test, BATCH_SIZE)

# ②	定义训练批次、周期、图像尺寸、学习率等超参数
# 已在文件开头定义


class Resnet34TransModel(pt.nn.Module):

    def __init__(self, n_cls, **kwargs):
        super().__init__(**kwargs)
        # ③	导入Resnet34主干网络作为预训练模型
        resnet = ptv.models.resnet34(pretrained=True)
        n_resnet_out_ch = resnet.fc.in_features
        self.resnet = pt.nn.Sequential(*(list(resnet.children())[:-1]))
        for p in self.resnet.parameters():
            p.requires_grad = False

        # ④	搭建后端网络模型
        self.fc = pt.nn.Linear(n_resnet_out_ch, n_cls)

    def forward(self, x):
        x = self.resnet(x)
        x = pt.squeeze(x, dim=3)
        x = pt.squeeze(x, dim=2)
        x = self.fc(x)
        return x


# ⑤	打印输出整体模型维度结构
model = Resnet34TransModel(n_cls).to(device)
print(model)

# ⑥	定义适当的损失函数loss和优化器optimizer
criterion = pt.nn.CrossEntropyLoss()
optim = pt.optim.Adam(model.parameters(), lr=ALPHA)


# ⑦	定义计算模型准确率函数
def accuracy(y_true, y_pred):
    y_true = y_true.long()
    y_pred = y_pred.argmax(dim=1).long()
    acc = pt.eq(y_true, y_pred).float().mean()
    return acc


# The procedure to process data
def process_data(dl, label, is_train, epoch=1, is_return_data=False):
    avg_loss, avg_acc = 0., 0.
    dl_len = len(dl)
    group = int(np.ceil(dl_len / 5))
    y_true, y_pred = None, None
    for i, (bx, by) in enumerate(dl):
        bx = bx.float().to(device)
        by = by.long().to(device)
        if is_train:
            model.train(True)
            optim.zero_grad()
            pred = model(bx)
            loss = criterion(pred, by)
            loss.backward()
            optim.step()
            model.train(False)
        else:
            model.train(False)
            pred = model(bx)
            loss = criterion(pred, by)
        acc = accuracy(by, pred)
        loss = loss.detach().cpu().numpy()
        acc = acc.detach().cpu().numpy()
        if i == dl_len - 1 or i % group == 0:
            print(f'{label}: epoch#{epoch}: batch#{i + 1}: loss = {loss}, acc = {acc}')
        avg_loss += loss
        avg_acc += acc
        if is_return_data:
            by = by.detach().cpu().numpy()
            pred = pred.detach().cpu().numpy().argmax(axis=1)
            by = np.asarray(by, dtype=np.int64)
            pred = np.asarray(pred, dtype=np.int64)
            if y_true is None:
                y_true = by
                y_pred = pred
            else:
                y_true = np.concatenate([y_true, by], axis=0)
                y_pred = np.concatenate([y_pred, pred], axis=0)
    avg_loss /= i + 1
    avg_acc /= i + 1
    if is_return_data:
        return avg_loss, avg_acc, y_true, y_pred
    else:
        return avg_loss, avg_acc

# ⑧	打印输出每一训练周期训练集和验证集的平均损失值和平均准确率
avg_loss_his, avg_acc_his = [], []
avg_loss_his_val, avg_acc_his_val = [], []
for epoch in range(N_EPOCHS):
    avg_loss, avg_acc = process_data(dl_train, 'train', True, epoch + 1)
    print(f'epoch#{epoch + 1}: avg_loss = {avg_loss}, avg_acc = {avg_acc}')
    avg_loss_val, avg_acc_val = process_data(dl_val, 'val', False, epoch + 1)
    print(f'epoch#{epoch + 1}: avg_loss = {avg_loss}, avg_acc = {avg_acc}, avg_loss_val = {avg_loss_val}, avg_acc_val = {avg_acc_val}')
    avg_loss_his.append(avg_loss)
    avg_loss_his_val.append(avg_loss_val)
    avg_acc_his.append(avg_acc)
    avg_acc_his_val.append(avg_acc_val)

# plotting parameters
spr = 2
spc = 2
spn = 0
plt.figure(figsize=[6, 6])

# ⑨	图示输出训练集和验证集的平均损失值曲线
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Loss')
plt.plot(avg_loss_his, label='train')
plt.plot(avg_loss_his_val, label='val')
plt.legend()
plt.grid()

# ⑩	图示输出训练集和验证集的平均准确率曲线
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Accuracy')
plt.plot(avg_acc_his, label='train')
plt.plot(avg_acc_his_val, label='val')
plt.legend()
plt.grid()

# 11	模型训练完成后，用测试集测试模型分类猫狗正确和不正确饼图
avg_loss_test, avg_acc_test, y_true, y_pred = process_data(dl_test, 'test', False, is_return_data=True)
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Accuracy on test')
plt.pie([avg_acc_test, 1 - avg_acc_test], explode=[0., 0.1], labels=['Correct', 'Wrong'], autopct='%0.2f%%')

# 12	用测试集数据，底层代码实现计算TN、TP、FN、FP的混淆矩阵，并热图输出
mat = confusion_matrix(y_true, y_pred)
print('Confusion matrix:')
print(mat)
spn += 1
plt.subplot(spr, spc, spn)
plt.title('Confusion matrix')
sns.heatmap(mat, annot=True)

# 13	公式计算输出查准率P和查全率R
print(f'Precision score = {precision_score(y_true, y_pred)}')
print(f'Recall score = {recall_score(y_true, y_pred)}')
# 14	公式计算输出准确率和F1分值
print(f'F1 score = {f1_score(y_true, y_pred)}')

# Finally show plotting
print('Please check and close the plotting window to finish ...')
plt.show()
print('Over.')
